import librosa
import matplotlib.pyplot as plt
import numpy as np
import sounddevice as sd
import soundfile as sf
from scipy.signal import resample
from ssqueezepy import ssq_stft, ssq_cwt

plugin_paths = {
    'Sylenth1': 'C:/VstPlugins/64bit/Sylenth1.dll',
    'Massive': 'C:/VstPlugins/64bit/Massive.dll',
}


def plot_spec(*specs, title=(), save_to=''):
    mx = max(len(specs), specs[0].shape[0])
    for i in range(len(specs)):
        for j in range(specs[i].shape[0]):
            plt.subplot2grid((mx, mx), (j, i), rowspan=(mx if specs[i].shape[0] == 1 else 1))
            plt.imshow(np.abs(specs[i][j, :, :, 0]), aspect='auto', cmap='jet')
            if j == 0:
                plt.title(title[i])

    if save_to:
        plt.savefig(save_to)
    plt.show()


def play(data, sr=44100, blocking=False, save_to=''):
    # print('正在播放：' + str(sr) + ' Hz')
    for i in range(data.shape[0]):
        if data.ndim == 3:
            if data.shape[1] == 1:
                dat = data[i, ...].reshape(-1)
            else:
                dat = librosa.to_mono(data[i, ...])
            sd.play(dat, sr, blocking=blocking)
            if save_to:
                sf.write(save_to, dat, sr)
        else:
            sd.play(data[i, ...], sr, blocking=blocking)
            if save_to:
                sf.write(save_to, data[i, ...], sr)
    # print('播放已结束')


def audio_to_spec(data, src_sr, cutoff_t=2., factor=2e-3):
    data_ = []
    data = data[:, :, :int(src_sr * cutoff_t)]

    for i in range(data.shape[0]):
        dat = librosa.to_mono(data[i, ...])
        dat = resample(dat, int(dat.shape[-1] * factor))
        dats = []
        dats += ssq_stft(dat)[:2]
        dats += ssq_cwt(dat)[:2]
        for j in range(len(dats)):
            dats[j] = dats[j].reshape(dats[j].shape + (1,))
        data_.append(np.concatenate(dats, axis=-1))

    return np.concatenate(data_).reshape((len(data_),) + data_[0].shape), src_sr * factor
